from audio_text_groq import call_groq_api
from audio_text_whisper import call_whisper_api
from audio_text_whisper import has_char_in_str
from common import convert_time


def gen_subtitle(
    video_path, output_path, split_chars="，。,.；;", type="local", model_name="medium"
):
    """
    使用 Whisper 模型识别视频中的语音，并生成 SRT 字幕文件

    Args:
        video_path: 视频文件路径
        output_path: 输出的 SRT 文件路径
    """
    if type == "groq":
        segments = call_groq_api(video_path, model_name)
    elif type == "local":
        segments = call_whisper_api(video_path, model_name)
    else:
        raise ValueError("Invalid type")
    out_list = []
    index = 1
    with open(output_path, "w", encoding="utf-8") as f:
        for segment in segments:
            seek = 0
            start_time = convert_time(segment["start"] + seek)
            end_time = convert_time(segment["end"] + seek)
            text = segment["text"].strip()

            if has_char_in_str(text, split_chars) and "words" in segment:
                item_start = None
                item_text = []
                for word_item in segment["words"]:
                    word = word_item["word"]
                    start = word_item["start"]
                    end = word_item["end"]
                    item_text.append(word)
                    if item_start is None:
                        item_start = start
                    if has_char_in_str(word, split_chars):
                        start_time = convert_time(item_start + seek)
                        end_time = convert_time(end + seek)
                        text = "".join(item_text)
                        f.write(f"{index}\n")
                        f.write(f"{start_time} --> {end_time}\n")
                        f.write(f"{text}\n\n")

                        out_list.append(f"{index}\n")
                        out_list.append(f"{start_time} --> {end_time}\n")
                        out_list.append(f"{text}\n\n")
                        index += 1
                        item_start = None
                        item_text = []
                if item_text:
                    start_time = convert_time(item_start + seek)
                    end_time = convert_time(end + seek)
                    text = "".join(item_text)
                    f.write(f"{index}\n")
                    f.write(f"{start_time} --> {end_time}\n")
                    f.write(f"{text}\n\n")
                    out_list.append(f"{index}\n")
                    out_list.append(f"{start_time} --> {end_time}\n")
                    out_list.append(f"{text}\n\n")
                    index += 1
            else:
                # 写入 SRT 格式
                print(f"{index} {start_time} {end_time} {text}")
                f.write(f"{index}\n")
                f.write(f"{start_time} --> {end_time}\n")
                f.write(f"{text}\n\n")

                out_list.append(f"{index}\n")
                out_list.append(f"{start_time} --> {end_time}\n")
                out_list.append(f"{text}\n\n")
                index += 1
    return "".join(out_list)
